/*****************************************************************
 Copyright (c) 2020, Unitree Robotics.Co.Ltd. All rights reserved.
******************************************************************/
#include "unitree_legged_sdk/unitree_legged_sdk.h"
#include <iostream>
#include <unistd.h>
#include <string.h>
#include <array>
#include <math.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <msgpack.hpp>

using namespace UNITREE_LEGGED_SDK;

class RobotInterface
{
public:
    RobotInterface(uint8_t level): safe(LeggedType::A1), udp(level){
        udp.InitCmdData(cmd);
    }
    void UDPRecv();
    void UDPSend();
    LowState getState();
    void RobotControl(int motor, std::array<int, 12> work,
                    //   std::array<uint8_t, 12> mode,
                      std::array<float, 12> q, std::array<float, 12> dq, std::array<float, 12> tau,
                      std::array<float, 12> Kp, std::array<float, 12> Kd);

    Safety safe;
    UDP udp;
    LowCmd cmd = {0};
    LowState state = {0};
    float dt = 0.002;     // 0.001~0.01
};


void RobotInterface::UDPRecv()
{
    udp.Recv();
}

void RobotInterface::UDPSend()
{  
    udp.Send();
}

LowState RobotInterface::getState()
{
    udp.GetRecv(state);
    return state;
}

void RobotInterface::RobotControl(int motor, std::array<int, 12> work,
                                //   std::array<uint8_t, 12> mode,
                                  std::array<float, 12> q, std::array<float, 12> dq, std::array<float, 12> tau,
                                  std::array<float, 12> Kp, std::array<float, 12> Kd) 
{
    if(motor>=0 && motor <=11)
    {
        // cmd.motorCmd[i].mode = 0;                // desired working mode
        cmd.motorCmd[motor].q = q[motor];           // desired angle (unit: radian) 
        cmd.motorCmd[motor].dq = dq[motor];         // desired velocity (unit: radian/second)
        cmd.motorCmd[motor].tau = tau[motor];       // desired output torque (unit: N.m)
        cmd.motorCmd[motor].Kp = Kp[motor];         // desired position stiffness (unit: N.m/rad )
        cmd.motorCmd[motor].Kd = Kd[motor];         // desired velocity stiffness (unit: N.m/(rad/s) )
    }
    else
    {
        for(int i = 0; i < 12; i++)
        {
            if(!work[i]) continue;
            // cmd.motorCmd[i].mode = 0;            // desired working mode
            cmd.motorCmd[i].q = q[i];               // desired angle (unit: radian) 
            cmd.motorCmd[i].dq = dq[i];             // desired velocity (unit: radian/second)
            cmd.motorCmd[i].tau = tau[i];           // desired output torque (unit: N.m)
            cmd.motorCmd[i].Kp = Kp[i];             // desired position stiffness (unit: N.m/rad )
            cmd.motorCmd[i].Kd = Kd[i];             // desired velocity stiffness (unit: N.m/(rad/s) )
            // for(int j = 0; j < 3; j++) cmd.motorCmd[i].reserve[j] = 0;
        }
    }
    
    udp.SetSend(cmd);
}

namespace py=pybind11;

PYBIND11_MODULE(robot_interface_low_level, m) 
{
    py::class_<RobotInterface>(m, "RobotInterface")
    .def(py::init<uint8_t>())
    .def("UDPRecv", &RobotInterface::UDPRecv)
    .def("UDPSend", &RobotInterface::UDPSend)
    .def("getState", &RobotInterface::getState)
    .def("robotControl", &RobotInterface::RobotControl);

    py::class_<Cartesian>(m, "Cartesian")
        .def(py::init<>())
        .def_readwrite("x", &Cartesian::x)
        .def_readwrite("y", &Cartesian::y)
        .def_readwrite("z", &Cartesian::z);

    py::class_<IMU>(m, "IMU")
        .def(py::init<>())
        .def_readwrite("quaternion", &IMU::quaternion)
        .def_readwrite("gyroscope", &IMU::gyroscope)
        .def_readwrite("accelerometer", &IMU::accelerometer)
        .def_readwrite("rpy", &IMU::rpy)
        .def_readwrite("temperature", &IMU::temperature);

    py::class_<LowState>(m, "LowState")
        .def(py::init<>())
        .def_readwrite("levelFlag", &LowState::levelFlag)
        .def_readwrite("commVersion", &LowState::commVersion)
        .def_readwrite("robotID", &LowState::robotID)
        .def_readwrite("SN", &LowState::SN)
        .def_readwrite("bandWidth", &LowState::bandWidth)
        .def_readwrite("imu", &LowState::imu)
        .def_readwrite("motorState", &LowState::motorState)
        .def_readwrite("footForce", &LowState::footForce)
        .def_readwrite("footForceEst", &LowState::footForceEst)
        .def_readwrite("tick", &LowState::tick)
        .def_readwrite("wirelessRemote", &LowState::wirelessRemote)
        .def_readwrite("reserve", &LowState::reserve)
        .def_readwrite("crc", &LowState::crc);
}
